#include "accumulator.h"

using namespace std;

extern double fx, fy, cx, cy;

void JacobianAccumulator::accumulate_jacobian(const cv::Range &range)
{
    // parameters
    const int half_patch_size = 1;
    int cnt_good = 0;
    Matrix6d hessian = Matrix6d::Zero();
    Vector6d bias = Vector6d::Zero();
    double cost_tmp = 0;

    for (size_t i = range.start; i < range.end; i++)
    {

        // compute the projection in the second image
        Eigen::Vector3d point_ref =
            depth_ref[i] * Eigen::Vector3d((px_ref[i][0] - cx) / fx, (px_ref[i][1] - cy) / fy, 1);
        Eigen::Vector3d point_cur = T21 * point_ref;
        if (point_cur[2] < 0) // depth invalid
            continue;

        float u = fx * point_cur[0] / point_cur[2] + cx, v = fy * point_cur[1] / point_cur[2] + cy;
        if (u < half_patch_size || u > img2.cols - half_patch_size || v < half_patch_size ||
            v > img2.rows - half_patch_size)
            continue;

        projection[i] = Eigen::Vector2d(u, v);
        double X = point_cur[0], Y = point_cur[1], Z = point_cur[2],
               Z2 = Z * Z, Z_inv = 1.0 / Z, Z2_inv = Z_inv * Z_inv;
        cnt_good++;

        // Projection equations:
        //           [u1]    1          [u2]    1              1
        //      p1 = |v1| = -- KP, p2 = [v2] = -- K(RP + t) = -- K(TP)
        //           [ 1]   Z1          [ 1]   Z2             Z2      (1:3)
        //
        //                   1
        //       q = TP, u = -- Kq
        //                   Z2
        // photometric error:
        //       e = I1(p1) - I2(p2) = I1(p1) - I2(u)
        // optimization:
        //       *                  N    T
        //      T  = argmin J(T) = SUM ei ei, ei = I1(p1,i) - I2(p2,i)
        //             T           i=1
        // gradients:
        //      ∂e   ∂I2 ∂u  ∂q
        //      -- = --- -- --- δξ 
        //      ∂T    ∂u ∂q ∂δξ
        //
        //     ∂I2   [∂I2 ∂I2]'
        //     --- = |--- ---|
        //      ∂u   [∂u2 ∂v2]
        // 
        //           [fx         fx*X       fx*XY         fx*X^2    fx*Y]
        //           |--    0  - ----    - -------   fx + ------  - ----|
        //      ∂u   | Z          Z^2        Z^2             Z^2      Z |   ∂u ∂q
        //     --- = |                                                  | = -- ---
        //     ∂δξ   |     fy    fy*Y        fy*Y^2     fy*XY      fy*X |   ∂q ∂δξ
        //           | 0   --  - ----  -fy - ------     -----      ---- |
        //           [      Z     Z^2          Z^2       Z^2         Z  ]
        //             ∂I2  ∂u
        //       J = - --- ---
        //              ∂u ∂δξ

        // compute error and jacobian
        for (int x = -half_patch_size; x <= half_patch_size; x++)
            for (int y = -half_patch_size; y <= half_patch_size; y++)
            {

                double error = GetGrayscale(img1, px_ref[i][0] + x, px_ref[i][1] + y) -
                               GetGrayscale(img2, u + x, v + y);
                Matrix26d J_pixel_xi;
                Eigen::Vector2d J_img_pixel;

                J_pixel_xi(0, 0) = fx * Z_inv;
                J_pixel_xi(0, 1) = 0;
                J_pixel_xi(0, 2) = -fx * X * Z2_inv;
                J_pixel_xi(0, 3) = -fx * X * Y * Z2_inv;
                J_pixel_xi(0, 4) = fx + fx * X * X * Z2_inv;
                J_pixel_xi(0, 5) = -fx * Y * Z_inv;

                J_pixel_xi(1, 0) = 0;
                J_pixel_xi(1, 1) = fy * Z_inv;
                J_pixel_xi(1, 2) = -fy * Y * Z2_inv;
                J_pixel_xi(1, 3) = -fy - fy * Y * Y * Z2_inv;
                J_pixel_xi(1, 4) = fy * X * Y * Z2_inv;
                J_pixel_xi(1, 5) = fy * X * Z_inv;

                J_img_pixel = Eigen::Vector2d(
                    0.5 * (GetGrayscale(img2, u + 1 + x, v + y) - GetGrayscale(img2, u - 1 + x, v + y)),
                    0.5 * (GetGrayscale(img2, u + x, v + 1 + y) - GetGrayscale(img2, u + x, v - 1 + y)));

                // total jacobian
                Vector6d J = -1.0 * (J_img_pixel.transpose() * J_pixel_xi).transpose();

                hessian += J * J.transpose();
                bias += -error * J;
                cost_tmp += error * error;
            }
    }

    if (cnt_good)
    {
        // set hessian, bias and cost
        unique_lock<mutex> lck(hessian_mutex);
        H += hessian;
        b += bias;
        cost += cost_tmp / cnt_good;
    }
}
